package pers.whj.util.graph;

import pers.whj.util.list.MinHeap;
import pers.whj.util.set.UFSet;

public class WeiGraphMatrix<T, E> extends AbstractGraph<T> implements WeiGraph<T, E>, WeiGraphOp<T, E> {
    protected Object[] mVertices;
    protected Object[][] mEdges;

    public WeiGraphMatrix() {
        super();
        mVertices = new Object[mMaxVertices];
        mEdges = new Object[mMaxVertices][mMaxVertices];
    }

    public WeiGraphMatrix(int maxVertices) {
        super(maxVertices);
        mVertices = new Object[mMaxVertices];
        mEdges = new Object[mMaxVertices][mMaxVertices];
    }

    @Override
    @SuppressWarnings({"unchecked"})
    protected Vertex<T> vertex(int index) {
        // assert index is valid
        return (Vertex<T>) mVertices[index];
    }

    @Override
    public boolean isCompleteGraph() {
        return mNumEdges == (long) mNumVertices / 2 * (mNumVertices - 1);
    }

    @Override
    public int getFirstNeighbor(int v) {
        checkIfVertexExists(v);
        for (int i = 0; i < mNumVertices; i++) {
            if (mEdges[v][i] != null) {
                return i;
            }
        }
        return -1;
    }

    @Override
    public int getNextNeighbor(int v, int w) {
        checkIfVerticesExist(v, w);
        for (int i = w + 1; i < mNumVertices; i++) {
            if (mEdges[v][i] != null) {
                return i;
            }
        }
        return -1;
    }

    @Override
    public int insertVertex(T vertex) {
        if (mNumVertices == mMaxVertices) {
            return -1;
        }
        mVertices[mNumVertices] = new Vertex<>(vertex);
        mNumVertices++;
        return mNumVertices - 1;
    }

    @Override
    public boolean insertEdge(int v1, int v2) {
        return false;
    }

    @Override
    public boolean insertEdge(T vertex1, T vertex2) {
        return false;
    }

    @Override
    public boolean removeVertex(int v) {
        checkIfVertexExists(v);
        mVertices[v] = mVertices[mNumVertices - 1];
        for (int i = 0; i < mNumVertices; i++) {
            if (mEdges[i][v] != null) {
                mNumEdges--;
            }
        }
        for (int i = 0; i < mNumVertices; i++) {
            mEdges[i][v] = mEdges[i][mNumVertices - 1];
        }
        mNumVertices--;
        System.arraycopy(mEdges[mNumVertices], 0, mEdges[v], 0, mNumVertices);
        return true;
    }

    @Override
    public boolean removeVertex(T vertex) {
        return removeVertex(getPosition(vertex));
    }

    @Override
    public boolean removeEdge(int v1, int v2) {
        checkIfVerticesExist(v1, v2);
        if (mEdges[v1][v2] != null) {
            mEdges[v1][v2] = mEdges[v2][v1] = null;
            mNumEdges--;
            return true;
        }
        return false;
    }

    @Override
    public boolean removeEdge(T vertex1, T vertex2) {
        return removeEdge(getPosition(vertex1), getPosition(vertex2));
    }

    @Override
    public boolean hasEdge(int v1, int v2) {
        checkIfVerticesExist(v1, v2);
        return mEdges[v1][v2] != null;
    }

    @Override
    public boolean hasEdge(T vertex1, T vertex2) {
        return hasEdge(getPosition(vertex1), getPosition(vertex2));
    }

    @Override
    public void removeAllEdges() {
        for (int i = 0; i < mNumVertices; i++) {
            for (int j = 0; j < mNumVertices; j++) {
                mEdges[i][j] = null;
            }
        }
        mNumEdges = 0;
    }

    @Override
    public void clear() {
        removeAllEdges();
        for (int i = 0; i < mNumVertices; i++) {
            mVertices[i] = null;
        }
        mNumVertices = 0;
    }

    @Override
    public boolean resize(int maxVertices) {
        if (maxVertices < mNumVertices) {
            return false;
        }
        if (maxVertices != mMaxVertices) {
            Object[] newVertices = new Object[maxVertices];
            System.arraycopy(mVertices, 0, newVertices, 0, mNumVertices);
            mVertices = newVertices;
            mMaxVertices = maxVertices;
        }
        return true;
    }

    @Override
    @SuppressWarnings({"unchecked"})
    public E getWeight(int v1, int v2) {
        checkIfVerticesExist(v1, v2);
        return (E) mEdges[v1][v2];
    }

    @Override
    public E getWeight(T vertex1, T vertex2) {
        return getWeight(getPosition(vertex1), getPosition(vertex2));
    }

    @Override
    public void setWeight(int v1, int v2, E weight) {
        checkIfVerticesExist(v1, v2);
        if (mEdges[v1][v2] != null && weight != null) {
            mEdges[v1][v2] = mEdges[v2][v1] = weight;
        }
    }

    @Override
    public void setWeight(T vertex1, T vertex2, E weight) {
        setWeight(getPosition(vertex1), getPosition(vertex2), weight);
    }

    @Override
    public boolean insertEdge(int v1, int v2, E weight) {
        checkIfVerticesExist(v1, v2);
        if (v1 == v2 || weight == null) {
            return false;
        }
        if (mEdges[v1][v2] != null) {
            return false;
        }
        mEdges[v1][v2] = mEdges[v2][v1] = weight;
        mNumEdges++;
        return true;
    }

    @Override
    public boolean insertEdge(T vertex1, T vertex2, E weight) {
        return insertEdge(getPosition(vertex1), getPosition(vertex2), weight);
    }

    @Override
    public void getMinSpanTreeByKruskal(MinSpanTree<E> minSpanTree) {
        if (minSpanTree != null && isConnectedGraph()) {
            MinHeap<MinSpanTree.Edge<E>> heap = new MinHeap<>();
            UFSet ufSets = new UFSet(mNumVertices);
            int i, j;
            for (i = 0; i < mNumVertices; i++) {
                for (j = i + 1; j < mNumVertices; j++) {
                    E weight = getWeight(i, j);
                    if (weight != null) {
                        heap.insert(new MinSpanTree.Edge<>(i, j, weight));
                    }
                }
            }
            int counts = 1;
            MinSpanTree.Edge<E> edge;
            while (counts < mNumVertices) {
                edge = heap.delete();
                i = ufSets.find(edge.head);
                j = ufSets.find(edge.tail);
                if (i != j) {
                    ufSets.weightedUnion(i, j);
                    minSpanTree.add(edge);
                    counts++;
                }
            }
        }
    }

    @Override
    public void getMinSpanTreeByPrim(MinSpanTree<E> minSpanTree, int v) {
        if (minSpanTree == null || !isConnectedGraph()) {
            return;
        }
        checkIfVertexExists(v);
        boolean[] vertices = new boolean[mNumVertices];
        vertices[v] = true;
        MinHeap<MinSpanTree.Edge<E>> heap = new MinHeap<>();
        int u;
        int counts = 1;
        do {
            u = getFirstNeighbor(v);
            while (u != -1) {
                if (!vertices[u]) {
                    heap.insert(new MinSpanTree.Edge<>(v, u, getWeight(v, u)));
                }
                u = getNextNeighbor(v, u);
            }
            while (!heap.isEmpty() && counts < mNumVertices) {
                MinSpanTree.Edge<E> edge = heap.delete();
                if (!vertices[edge.tail]) {
                    minSpanTree.add(edge);
                    v = edge.tail;
                    vertices[v] = true;
                    counts++;
                    break;
                }
            }
        } while (counts < mNumVertices);
    }

    @Override
    public void getMinSpanTreeByPrim(MinSpanTree<E> minSpanTree, T vertex) {
        getMinSpanTreeByPrim(minSpanTree, getPosition(vertex));
    }
}